import numpy as np
from scipy.signal import spectrogram
import pywt as wt
from scipy.misc import imresize
from sklearn.preprocessing import normalize


class fp_obj(object):
    def __init__(self, sampling_rate, window_length, window_lag, fingerprint_length, fingerprint_lag,
                 min_freq=0, max_freq=None, nfreq=32, ntimes=64):
        self.sampling_rate = sampling_rate
        self.window_len = window_length
        self.window_lag = window_lag
        self.fp_len = fingerprint_length
        self.fp_lag = fingerprint_lag
        self.max_freq = self._initialize_frequencies(max_freq)
        self.min_freq = min_freq
        self.new_d1 = int(nfreq)
        self.new_d2 = int(ntimes)
        self.d1 = None
        self.d2 = None
        self.haar_means = None
        self.haar_stddevs = None
        self.haar_medians = None
        self.haar_absdevs = None

    def get_window_params(self, N, L, dL):
        idx0 = np.asarray(range(0, N + 1, dL))
        idx2 = np.asarray(range(L, N + 1, dL))
        nWindows = len(idx2)
        idx1 = idx0[0:nWindows]
        print(idx1[:100])
        print(idx2[:100])
        return nWindows, idx1, idx2

    # 短时傅里叶变换得到频谱图
    def data_to_spectrogram(self, x_data, window_type='hanning'):
        print('data_to_spectrogram')
        f, t, Sxx = spectrogram(x_data, fs=self.sampling_rate, window=window_type,
                                nperseg=int(self.sampling_rate * self.window_len),
                                noverlap=int(self.sampling_rate * (self.window_len - self.window_lag)))
        # 带通滤波
        if self.min_freq > 0:
            fidx_keep = (f >= self.min_freq)
            Sxx = Sxx[fidx_keep, :]
            f = f[fidx_keep]

        if self.max_freq < f[-1]:
            fidx_keep = (f <= self.max_freq)
            Sxx = Sxx[fidx_keep, :]
            f = f[fidx_keep]
        self.frequencies = f
        self.times = t
        return f, t, Sxx

    # 为了检测出短时间的事件，分割频谱图
    def spectrogram_to_spectral_images(self, Sxx):
        print('spectrogram_to_spectral_images')
        nFreq, nTimes = np.shape(Sxx)  # (频率：61, 时间：373996) 原数据持续时间为74805.05
        print(np.shape(Sxx))
        print('nFreq', nFreq)
        print('nTimes', nTimes)
        print('self.fp_len', self.fp_len)
        print('self.fp_lag', self.fp_lag)
        nWindows, idx1, idx2 = self.get_window_params(nTimes, self.fp_len, self.fp_lag)
        print('nWindows', nWindows)
        spectral_images = np.zeros([nWindows, nFreq, self.fp_len])
        for i in range(nWindows):
            # print('i :',i)
            spectral_images[i, :, :] = Sxx[:, idx1[i]:idx2[i]]
            # print('spectral_images[i, :, :]',spectral_images[i, :, :])
        self.nWindows = nWindows
        nWindows, self.d1, self.d2 = np.shape(spectral_images)
        return spectral_images, nWindows, idx1, idx2

    # 对不符合要求的频谱图重新调整大小
    def _resize_spectral_images(self, spectral_images, new_d1, new_d2):
        new_spectral_images = np.zeros([self.nWindows, new_d1, new_d2])
        for i in range(self.nWindows):
            new_spectral_images[i, :, :] = imresize(spectral_images[i, :, :], (new_d1, new_d2), interp='bilinear',
                                                    mode='F')
        return new_spectral_images

    # 变形(不懂)
    def _unwrap_wavelet_coeffs(self, coeffs):
        # print('_unwrap_wavelet_coeffs')
        L = len(coeffs)
        cA = coeffs[0]
        for i in range(1, L):
            (cH, cV, cD) = coeffs[i]
            cA = np.concatenate((np.concatenate((cA, cV), axis=1), np.concatenate((cH, cD), axis=1)), axis=0)
        # print('cA', np.shape(cA))
        return cA

    # 对于分割后的每一个频谱图，计算小波变换级数
    def spectral_images_to_wavelet(self, spectral_images, wavelet=wt.Wavelet('db1')):
        print('spectral_images_to_wavelet')
        if (int(self.new_d1) != self.d1) or (int(self.new_d2) != self.d2):
            spectral_images = self._resize_spectral_images(spectral_images, self.new_d1, self.new_d2)
        haar_images = np.zeros([self.nWindows, self.new_d1, self.new_d2])
        print('self.nWindows', self.nWindows)
        for i in range(self.nWindows):
            # print('spectral_images[i, :, :]',spectral_images[i, :, :])
            # print('正在变换第%s个窗口' % i)
            coeffs = wt.wavedec2(spectral_images[i, :, :], wavelet)
            haar_images[i, :, :] = self._unwrap_wavelet_coeffs(coeffs)
        return haar_images

    def _initialize_frequencies(self, max_freq):
        if max_freq is None:
            max_freq = self.sampling_rate / 2.0
        return max_freq

    def data_to_haar_images(self, x_data):
        print('data_to_haar_images')
        f, t, Sxx = self.data_to_spectrogram(x_data)
        print('np.shape(Sxx) : ', np.shape(Sxx))
        spectral_images, nWindows, idx1, idx2 = self.spectrogram_to_spectral_images(Sxx)
        haar_images = self.spectral_images_to_wavelet(spectral_images)
        print('小波变换完毕,小波频谱如下')
        print(np.shape(haar_images))  # (74793, 1024)
        haar_images = normalize(self._images_to_vectors(haar_images), axis=1)
        print('输出结果：', np.shape(haar_images))  # (74793, 1024)
        return haar_images, nWindows, idx1, idx2, Sxx, t

    # / 频谱转为向量
    def _images_to_vectors(self, images):
        print('_images_to_vectors')
        N, d1, d2 = np.shape(images)
        vectors = np.zeros([N, d1 * d2])
        for i in range(N):
            vectors[i, :] = np.reshape(images[i, :, :], (1, d1 * d2))
        return vectors

    def compute_haar_stats(self, haar_images, type=None):
        if type is 'MAD':
            shape = haar_images.shape
            print('shape', shape)
            medians = []
            for i in range(shape[1]):
                medians.append(np.median(haar_images[:, i]))
            self.haar_medians = np.array(medians)

            mad = []
            for i in range(shape[1]):
                tmp = abs(haar_images[:, i] - medians[i])
                mad.append(np.median(tmp))
            self.haar_absdevs = np.array(mad)

            return self.haar_medians, self.haar_absdevs
        if type is 'Zscore':
            self.haar_means = np.mean(haar_images, axis=0)
            self.haar_stddevs = np.std(haar_images, axis=0)
            return self.haar_means, self.haar_stddevs

    # fingerprint阶段函数
    # 标准化
    def standardize_haar(self, haar_images, type='MAD'):
        if type is 'Zscore':
            haar_images = (haar_images - self.haar_means) / self.haar_stddevs
            return haar_images
        elif type is 'MAD':
            haar_images = (haar_images - self.haar_medians) / self.haar_absdevs
            return haar_images
        else:
            print('Warning: invalid type - select type MAD or Zscore')
            return None

    # 需要认真研究
    def binarize_vectors_topK_sign(self, coeff_vectors, K):
        self.K = K
        N, M = np.shape(coeff_vectors)
        binary_vectors = np.zeros((N, 2 * M), dtype=bool)
        for i in range(N):
            idx = np.argsort(abs(coeff_vectors[i, :]))[-K:]
            binary_vectors[i, idx] = coeff_vectors[i, idx] > 0
            binary_vectors[i, idx + M] = coeff_vectors[i, idx] < 0
        return binary_vectors
